import fitz
from PIL import Image
import pytesseract
from docx import Document
from pptx import Presentation
import io
from typing import List, Union

async def perform_ocr(content: bytes, filename: str, lang: str) -> str:
    """执行OCR识别"""
    try:
        if not content:
            raise ValueError("文件内容为空")

        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            image = Image.open(io.BytesIO(content))
            if image.mode != 'RGB':
                image = image.convert('RGB')
            return pytesseract.image_to_string(image, lang=lang)
        else:
            pdf_document = fitz.open(stream=content, filetype="pdf")
            if pdf_document.page_count == 0:
                raise ValueError("PDF文件为空")

            text = ""
            for page in pdf_document:
                pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                text += pytesseract.image_to_string(img, lang=lang) + "\n"
            return text.strip()
    except Exception as e:
        raise Exception(f"OCR处理失败: {str(e)}")

async def pdf_to_word(content: bytes) -> bytes:
    """PDF转Word"""
    try:
        doc = Document()
        pdf_document = fitz.open(stream=content, filetype="pdf")
        
        for page in pdf_document:
            text = page.get_text()
            doc.add_paragraph(text)
        
        docx_bytes = io.BytesIO()
        doc.save(docx_bytes)
        return docx_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF转Word失败: {str(e)}")

async def pdf_to_ppt(content: bytes) -> bytes:
    """PDF转PPT"""
    try:
        prs = Presentation()
        pdf_document = fitz.open(stream=content, filetype="pdf")
        
        for page in pdf_document:
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            
            slide = prs.slides.add_slide(prs.slide_layouts[6])
            
            img_bytes = io.BytesIO()
            img.save(img_bytes, format='PNG')
            img_bytes.seek(0)
            
            slide.shapes.add_picture(
                img_bytes, 
                0, 0,
                width=prs.slide_width,
                height=prs.slide_height
            )
        
        pptx_bytes = io.BytesIO()
        prs.save(pptx_bytes)
        return pptx_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF转PPT失败: {str(e)}")

async def images_to_pdf(contents: List[bytes]) -> bytes:
    """图片转PDF"""
    try:
        pdf = fitz.open()
        
        for content in contents:
            image = Image.open(io.BytesIO(content))
            img_bytes = io.BytesIO()
            image.save(img_bytes, format='PNG')
            img_bytes.seek(0)
            
            img_doc = fitz.open("png", img_bytes.read())
            pdf.insert_pdf(img_doc)
        
        pdf_bytes = io.BytesIO()
        pdf.save(pdf_bytes)
        return pdf_bytes.getvalue()
    except Exception as e:
        raise Exception(f"图片转PDF失败: {str(e)}")

async def merge_pdfs(contents: List[bytes]) -> bytes:
    """合并PDF"""
    try:
        merged_pdf = fitz.open()
        
        for content in contents:
            pdf = fitz.open(stream=content, filetype="pdf")
            merged_pdf.insert_pdf(pdf)
        
        pdf_bytes = io.BytesIO()
        merged_pdf.save(pdf_bytes)
        return pdf_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF合并失败: {str(e)}")

async def split_pdf(content: bytes, page_ranges: str) -> bytes:
    """拆分PDF"""
    try:
        pdf = fitz.open(stream=content, filetype="pdf")
        new_pdf = fitz.open()
        
        # 解析页码范围
        pages = []
        for part in page_ranges.split(','):
            if '-' in part:
                start, end = map(int, part.split('-'))
                pages.extend(range(start-1, end))
            else:
                page_num = int(part)
                pages.append(page_num-1)
        
        # 提取页面
        for page_num in pages:
            if 0 <= page_num < len(pdf):
                new_pdf.insert_pdf(pdf, from_page=page_num, to_page=page_num)
        
        pdf_bytes = io.BytesIO()
        new_pdf.save(pdf_bytes)
        return pdf_bytes.getvalue()
    except Exception as e:
        raise Exception(f"PDF拆分失败: {str(e)}") 